/**
 * @file example_utils.cpp
 * @brief utils for convolution fprop
 * @version 0.1
 * @date 2022-06-22
 *
 * @copyright Copyright (c) 2021
 *
 */
#include "example_utils.h"

#include <math.h>
#include <stdio.h>

void FwdConvCpu(float *dst,
                const float *src,
                const float *fil,
                uint32_t bs,
                uint32_t fu,
                uint32_t fv,
                uint32_t inx,
                uint32_t iny,
                uint32_t inc,
                uint32_t onx,
                uint32_t ony,
                uint32_t onc) {
    for (uint32_t bt = 0; bt < bs; ++bt) {
        for (uint32_t oc = 0; oc < onc; ++oc) {
            for (uint32_t y = 0; y < ony; ++y) {
                for (uint32_t x = 0; x < onx; ++x) {
                    float s = 0.f;
                    for (uint32_t ic = 0; ic < inc; ++ic) {
                        for (uint32_t v = 0; v < fv; ++v) {
                            for (uint32_t u = 0; u < fu; ++u) {
                                s += src[((bt * inc + ic) * iny + (y + v)) * inx + x + u] *
                                     fil[((oc * inc + ic) * fv + v) * fu + u];
                            }
                        }
                    }
                    dst[((bt * onc + oc) * ony + y) * onx + x] = s;
                }
            }
        }
    }
}

bool Check(const float *a, const float *b, uint32_t n, double eps, uint32_t skip) {
    double s0        = 0;
    double s1        = 0;
    double precision = 1e-6;

    for (uint32_t i = 0; i < n; ++i) {
        double ai = static_cast<double>(a[i]);
        double bi = static_cast<double>(b[i]);
        double d  = ai - bi;
        double dd = d * d;

        if (skip > 0 && dd > precision) {
            skip--;
            continue;
        }

        double aa = 2.0 * ai * ai;
        s0 += dd;
        s1 += aa;
    }
    double rel_err = sqrt(s0 / (s1 + precision));
    if (rel_err > eps) {
        printf("\033[0;31mCheck fail: relative_diff %e, eps %e\033[0m\n", rel_err, eps);
    } else {
        printf("\033[0;32mCheck ok: relative_diff %e, eps %e\033[0m\n", rel_err, eps);
    }
    return (rel_err < eps);
}

void PaddingCPU(float *dst,
                const float *src,
                const uint32_t dst_elements,
                int *pad,
                uint32_t h,
                uint32_t w,
                uint32_t batch,
                uint32_t in_channel) {
    memset(dst, 0, sizeof(float) * dst_elements);
    uint32_t padding_h = h + pad[0] + pad[1];
    uint32_t padding_w = w + pad[2] + pad[3];
    for (int i1 = 0; i1 < batch; ++i1) {
        for (int i2 = 0; i2 < in_channel; ++i2) {
            for (int i3 = 0; i3 < h; ++i3) {
                for (int i4 = 0; i4 < w; ++i4) {
                    dst[((i1 * in_channel + i2) * padding_h + i3 + pad[0]) * padding_w + i4 +
                        pad[2]] = src[((i1 * in_channel + i2) * h + i3) * w + i4];
                }
            }
        }
    }
}

void StripArgs(char *s) {
    size_t len    = strlen(s);
    size_t offset = 0;
    for (int i = 0; i < len; ++i) {
        char c = s[i];
        if (c == '\t' || c == '\n' || c == 0x0d || c == 0x0a) {
            ++offset;
        } else {
            s[i - offset] = c;
        }
    }
    s[len - offset] = '\0';
}

bool FindArg(int argc, char *argv[], char *arg) {
    for (int i = 0; i < argc; ++i) {
        if (!argv[i]) continue;
        if (0 == strcmp(argv[i], arg)) {
            return 1;
        }
    }
    return 0;
}

void DelArg(int argc, char **argv, int index) {
    int i;
    for (i = index; i < argc - 1; ++i) {
        argv[i] = argv[i + 1];
    }
    argv[i] = 0;
}

int FindIntArg(int argc, char **argv, const char *arg, int def) {
    for (int i = 0; i < argc - 1; ++i) {
        if (!argv[i]) continue;
        if (0 == strcmp(argv[i], arg)) {
            def = atoi(argv[i + 1]);
            DelArg(argc, argv, i);
            DelArg(argc, argv, i);
            break;
        }
    }
    return def;
}

float FindFloatArg(int argc, char **argv, const char *arg, float def) {
    for (int i = 0; i < argc - 1; ++i) {
        if (!argv[i]) continue;
        if (0 == strcmp(argv[i], arg)) {
            def = atof(argv[i + 1]);
            DelArg(argc, argv, i);
            DelArg(argc, argv, i);
            break;
        }
    }
    return def;
}
